import sys
import importlib.util
import pandas as pd
import numpy as np
from pathlib import Path
from datetime import datetime
from zipline import run_algorithm
from zipline.api import symbol, record
from utils.datasource import GetHistoryData

def load_strategy_module(file_path):
    """动态加载策略模块"""
    module_name = file_path.stem
    spec = importlib.util.spec_from_file_location(module_name, file_path)
    if spec is None:
        raise ImportError(f"无法加载策略模块: {file_path}")
    module = importlib.util.module_from_spec(spec)
    if spec.loader is None:
        raise ImportError(f"策略模块加载器不可用: {file_path}")
    spec.loader.exec_module(module)
    return module

def run_backtest(strategy_file, start_date, end_date):
    """运行回测并收集结果（支持多资产组合）"""
    # 加载策略模块
    strategy_module = load_strategy_module(strategy_file)
    
    # 创建策略上下文对象
    class Context:
        def __init__(self):
            self.security_list = []
            self.index_components = []
            self.portfolio = type('Portfolio', (object,), {'positions': {}})()
    
    context = Context()
    
    # 模拟initialize调用
    strategy_module.initialize(context)
    
    # 获取CSI 300成分股
    from strategies.utils.a_share_utils import get_csi300_constituents
    context.index_components = get_csi300_constituents(context, data=None)
    
    # 获取历史数据（全股票池）
    all_data = {}
    for symbol in context.index_components:
        stock_data = GetHistoryData(symbol, adjust='hfq')
        stock_data = stock_data.reset_index()
        stock_data['date'] = pd.to_datetime(stock_data['date'])
        stock_data = stock_data.set_index('date')
        stock_data = stock_data.rename(columns={
            'open': 'open',
            'high': 'high',
            'low': 'low',
            'close': 'close',
            'volume': 'volume'
        })
        stock_data['symbol'] = symbol
        all_data[symbol] = stock_data
    
    # 创建组合数据
    portfolio_data = pd.concat(all_data.values())
    
    # 创建临时目录
    import os
    import shutil
    temp_dir = '/tmp/zipline'
    if os.path.exists(temp_dir):
        shutil.rmtree(temp_dir)
    os.makedirs(temp_dir)
    
    # 保存数据到CSV
    for symbol, df in all_data.items():
        df.to_csv(f"{temp_dir}/{symbol}.csv")
    
    # 创建Zipline数据包
    from zipline.data.bundles import register
    from zipline.data.bundles.csvdir import csvdir_equities
    
    # 注册临时数据包
    register(
        'a_share_bundle',
        csvdir_equities(
            ['daily'],
            temp_dir
        ),
    )
    
    # 运行回测
    try:
        results = run_algorithm(
            start=pd.Timestamp(start_date, tz='utc'),
            end=pd.Timestamp(end_date, tz='utc'),
            initialize=strategy_module.initialize,
            handle_data=strategy_module.handle_data,
            capital_base=1000000,
            data_frequency='daily',
            bundle='a_share_bundle'
        )
    except Exception as e:
        print(f"回测执行失败: {e}")
        return None
        
    if results is None or results.empty:
        print("回测未返回结果")
        return None
    
    # 收集A股特有机制验证数据
    a_share_metrics = {
        'csi300_coverage': len(context.index_components),
        'price_limit_events': 0,  # 将在策略中记录
        't1_violations': 0        # 将在策略中记录
    }
    
    # 计算关键指标
    if 'returns' not in results.columns:
        print("回测结果中缺少收益率数据")
        return None
        
    returns = results['returns']
    if len(returns) == 0:
        print("回测未生成收益率数据")
        return None
        
    cumulative_returns = (1 + returns).cumprod() - 1
    max_drawdown = (cumulative_returns / cumulative_returns.cummax() - 1).min()
    annual_return = (1 + cumulative_returns.iloc[-1]) ** (252/len(returns)) - 1
    sharpe_ratio = (annual_return - 0.03) / (returns.std() * np.sqrt(252)) if returns.std() > 0 else 0
    
    # 收集风险控制数据
    risk_controls = {
        'stop_loss_triggered': results.get('stop_loss_events', 0),
        'position_changes': len(results.orders) if 'orders' in results.columns else 0
    }
    
    return {
        'annual_return': annual_return,
        'max_drawdown': max_drawdown,
        'sharpe_ratio': sharpe_ratio,
        'risk_controls': risk_controls,
        'a_share_metrics': a_share_metrics,
        'results': results
    }

if __name__ == "__main__":
    if len(sys.argv) != 2:
        print("Usage: python run_backtest.py <strategy_file>")
        sys.exit(1)
    
    strategy_file = Path(sys.argv[1])
    start_date = '2018-01-01'
    end_date = '2023-12-31'
    
    results = run_backtest(strategy_file, start_date, end_date)
    
    # 输出结果
    if results:
        print(f"策略: {strategy_file.stem}")
        print(f"年化收益率: {results['annual_return']:.2%}")
        print(f"最大回撤: {results['max_drawdown']:.2%}")
        print(f"夏普比率: {results['sharpe_ratio']:.2f}")
        print(f"止损触发次数: {results['risk_controls']['stop_loss_triggered']}")
        print(f"仓位调整次数: {results['risk_controls']['position_changes']}")
    else:
        print("回测未生成有效结果")